import numpy as np

def point_to_point_distance(x_a, x_b):
    return np.linalg.norm(x_a - x_b)

def point_to_edge_distance(x_a, x_b, x_c):
    return np.linalg.norm(np.cross(x_a - x_b, x_b - x_c)) / np.linalg.norm(x_a - x_b)

def edge_to_edge_distance(x_i, x_i1, x_j, x_j1):
    u = np.cross(x_i1 - x_i, x_j1 - x_j)
    return np.abs(np.dot((x_i - x_j), u / np.linalg.norm(u)))

def calculate_min_distances_for_each_node(rod_nodes):
    num_rods = len(rod_nodes)
    min_distances = []

    for i in range(num_rods):
        for k in range(len(rod_nodes[i])):
            min_distances_for_rod = []
            for j in range(num_rods):
                if i == j:
                    continue
                min_distance = float('inf')

                # Point-to-point distance
                for l in range(len(rod_nodes[j])):
                    x_i = rod_nodes[i][k]
                    x_j = rod_nodes[j][l]
                    dist = point_to_point_distance(x_i, x_j)
                    if dist < min_distance:
                        min_distance = dist

                # Point-to-edge distance
                for l in range(len(rod_nodes[j]) - 1):
                    x_i = rod_nodes[i][k]
                    x_j, x_j1 = rod_nodes[j][l], rod_nodes[j][l + 1]
                    dist = point_to_edge_distance(x_i, x_j, x_j1)
                    if dist < min_distance:
                        min_distance = dist

                # Edge-to-edge distance (only if k < len(rod_nodes[i]) - 1)
                if k < len(rod_nodes[i]) - 1:
                    x_i, x_i1 = rod_nodes[i][k], rod_nodes[i][k + 1]
                    for l in range(len(rod_nodes[j]) - 1):
                        x_j, x_j1 = rod_nodes[j][l], rod_nodes[j][l + 1]
                        dist = edge_to_edge_distance(x_i, x_i1, x_j, x_j1)
                        if dist < min_distance:
                            min_distance = dist

                min_distances_for_rod.append((i, k, j, min_distance))
            min_distances.extend(min_distances_for_rod)

    return min_distances

def calculate_contact_energy(min_distance, h=0.01, delta=0.005):
    k = 15 / delta
    if min_distance > 2 * h + delta:
        return 0
    elif min_distance <= 2 * h - delta:
        return k * (2 * h - min_distance) ** 2
    elif min_distance <= 2 * h + delta and min_distance > 2 * h - delta:
        return (1 / k) * np.log(1 + np.exp(k * (2 * h + delta - min_distance))) ** 2
    else:
        return float('inf')

def calculate_contact_force(min_distance, h=0.01, delta=0.005):
    k = 15 / delta
    if min_distance > 2 * h + delta:
        return 0
    elif min_distance <= 2 * h - delta:
        return -2 * k * (2 * h - min_distance)
    elif min_distance <= 2 * h + delta and min_distance > 2 * h - delta:
        exp_term = np.exp(k * (2 * h + delta - min_distance))
        return -2 * (1 / k) * (k * exp_term / (1 + exp_term))
    else:
        return float('inf')

def calculate_friction_force(contact_force, mu=0.3, gamma=1.0, v_t_rel=np.array([1, 0, 0])):
    # Friction force includes the friction coefficient, smoothing factor, and tangential velocity direction
    friction_force = -mu * gamma * v_t_rel * abs(contact_force)
    return friction_force

# Example data for three rods with 10 nodes each
np.random.seed(0)  # Set seed for reproducibility
rod1 = np.random.rand(10, 3)
rod2 = np.random.rand(10, 3)
rod3 = np.random.rand(10, 3)

rods = [rod1, rod2, rod3]

# Calculate the minimum distances between each node and all other rods
min_distances = calculate_min_distances_for_each_node(rods)

# Calculate and display contact energies, forces, and friction forces
for (i, k, j, min_dist) in min_distances:
    contact_energy = calculate_contact_energy(min_dist)
    contact_force = calculate_contact_force(min_dist)
    v_t_rel = np.array([1, 0, 0])  # Example tangential relative velocity direction (unit vector)
    gamma = 1.0  # Smoothing factor, can be modified as needed
    friction_force = calculate_friction_force(contact_force, mu=0.3, gamma=gamma, v_t_rel=v_t_rel)
    print(f"Rod {i} Node {k} Minimum Distance to Rod {j}: {min_dist:.4f}, Contact Energy: {contact_energy:.4f}, Contact Force: {contact_force:.4f}, Friction Force: {friction_force}")
